#!/usr/bin/python
#Author: Jimmy Saw
#Usage: python get_lineage.py blastfile.tbl 50
#50 is max number of top hits to return
#It is supposed to give you the parent phylum of the hit organism from a BLAST result
#Needs to be in tabular format parsed by Blast2table Perl script

import sys
import re
from Bio import Entrez

def fmt(f):
    st = '{0:.4}'.format(f)
    return st

Entrez.email = "jimmysaw@gmail.com"

blastfile = sys.argv[1]
bf = open(blastfile, "rU")
bfl = bf.readlines()

si = blastfile.split('.')
seqid = si[0]

maxorgs = sys.argv[2]

#hit1 = re.compile('\[(.*)\]')
#hit2 = re.compile('\[(\w+ \w+ \w+)\].*')
#hit3 = re.compile('\[(\w+ \w+ \w+ \w+)\].*')
hit1 = re.compile('\[(.*)\]')
hit2 = re.compile('\[(.*)\] .*\[(.*)\]')
hit3 = re.compile('\[(.*)\] .*\[(.*)\] .*\[(.*)\]')

lin = re.compile('(\w+ \w+).*')
#lin = re.compile('(\w+).*')
acc = re.compile('\w+|(.*)|')

num = len(bfl)

maximum = 0

if num > int(maxorgs):
    maximum = int(maxorgs)
else:
    maximum = num

phyla = []
if len(bfl) == 0:
    #print blastfile + " has no BLAST hit."
    pass
else:
    for x, line in enumerate(bfl[0:maximum]):
        s = line.split('\t')
        evalue = s[1]
        identities = s[3]
        similarities = s[4]
        locus_tag = s[5]
        accession = s[9]
        desc = s[13].rstrip()
        searchterm = ""
        if hit3.search(desc):
            p = hit3.search(desc)
            searchterm = p.group(1)
        elif hit2.search(desc):
            p = hit2.search(desc)
            searchterm = p.group(1)
        elif hit1.search(desc):
            p = hit1.search(desc)
            searchterm = p.group(1)
        else:
            break

        handle = Entrez.esearch(db="Taxonomy", term=searchterm)
        record = Entrez.read(handle)
        taxid = record["IdList"]
        if len(taxid) != 0:
            h = Entrez.efetch(db="Taxonomy", id=taxid, retmode="xml")
            rec = Entrez.parse(h)      
            for r in rec:
                lineage = r.get("Lineage")
                l = lineage.split(';')
                #print l[3].lstrip()
                #print "Search term used =", searchterm #pattern.group(1)
                #print "Lineage found =", lineage
                #print x, "Phylum =", l[2].strip()
                phyla.append(l[2].strip())
        else:
            #print x, "Can't query NCBI Taxonomy. Taxid not found!"
            #print "Search term used =", searchterm
            pass

pls = len(phyla)
if pls > 0:
    unique_phyla = set(phyla)
    unique_phyla_list = list(unique_phyla)
    descriptions = []
    name_votes = []
    for d in unique_phyla_list:
        descriptions.append(d)
        count = phyla.count(d)
        name_votes.append(count)
    smallest = min(name_votes)
    largest = max(name_votes)
    max_index = name_votes.index(largest)
    min_index = name_votes.index(smallest)
    print seqid + "\t" + "Largest vote = " + str(largest) + "/" + str(maximum)\
        + " (" + str(fmt((float(largest)/float(maximum))*100)) + "%)" + "\t" + descriptions[max_index]
    #print "Smallest vote = " + str(smallest) + "\t" + descriptions[min_index]
else:
    print seqid + "\t" + "No BLAST hit or can't query taxid" + "\t" + "-"

bf.close()
